package com.seanLab.tool.TagSuggestModel.ExpandRank;

import com.seanLab.tool.TagSuggestModel.StaticModelProperties;

import java.io.BufferedReader;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.HashSet;
import java.util.logging.Logger;

/**
 * 根据公司词典调整TextRank最终返回分数（adjustRank），对于公司名给予1.5倍的分数
 */
public class EconomyTextRankKE extends FakeExpandRankKE {
    private static Logger LOG = Logger.getLogger(EconomyTextRankKE.class.getName());
    private static HashSet<String> company = null;

    private void loadDict(String path) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(new FileInputStream(path)));
        String line;
        while ((line = br.readLine()) != null) {
            String[] words = line.split("\t");
            if (words.length > 0) {
                company.add(words[0]);
            }
        }
        br.close();
    }

    public void loadModel(String modelPath) throws IOException {
        if (company == null) {
            company = new HashSet<>();
            // add all company dict here
            for (String dic : StaticModelProperties.companyDics) {
                LOG.info("Economy text rank loading " + dic);
                loadDict(dic);
            }
//            loadDict("workingDir/dict/das_sogou_comp.dic");
//            loadDict("workingDir/dict/das_stocks.dic");
            LOG.info("Economy text rank loaded! Size = " + company.size());
        }
    }

    public double adjustRank(String word, double score) {
        if (company.contains(word)) {
            return score * 1.5;
        }
        return score;
    }
}
